"""
Dice系数 / F1 Score 计算
"""
import torch
import numpy as np


def calculate_dice(pred, target, threshold=0.5, smooth=1e-6):
    """
    计算Dice系数（也称为F1 Score）

    Dice = 2 * |X ∩ Y| / (|X| + |Y|)

    Args:
        pred: 预测结果 (B, 1, H, W) 或 (B, C, H, W)
        target: 真实标签 (B, 1, H, W)
        threshold: 二值化阈值
        smooth: 平滑项，避免除零

    Returns:
        dice: Dice系数
    """
    if isinstance(pred, torch.Tensor):
        pred = pred.detach().cpu().numpy()
    if isinstance(target, torch.Tensor):
        target = target.detach().cpu().numpy()

    # 如果是多通道输出，取前景通道
    if pred.shape[1] > 1:
        pred = pred[:, 1:2, :, :]

    # 展平
    pred = pred.flatten()
    target = target.flatten()

    # 二值化
    pred = (pred > threshold).astype(np.float32)

    # 计算交集
    intersection = np.sum(pred * target)

    # 计算Dice
    dice = (2.0 * intersection + smooth) / (np.sum(pred) + np.sum(target) + smooth)

    return dice
